import tensorflow as tf
import numpy as np
import pickle
from tensorflow.keras import models, layers

with open("data_set.pickle", 'rb') as f:
    data_set = pickle.load(f)
x_train = data_set["x_train"]
y_train = data_set["y_train"]
x_test = data_set["x_test"]
y_test = data_set["y_test"]
x_train = np.reshape(x_train, [-1, 56, 56, 1])
x_test = np.reshape(x_test, [-1, 56, 56, 1])
print(x_train.shape)
print(y_train.shape)
print(x_test.shape)
print(y_test.shape)
model = models.Sequential()
# 第1层卷积，卷积核大小为3*3，32个，28*28为待训练图片的大小
model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=[56, 56, 1]))
model.add(layers.MaxPooling2D((2, 2)))
# 第2层卷积，卷积核大小为3*3，64个
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
# 第3层卷积，卷积核大小为3*3，64个
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.Flatten())
model.add(layers.Dense(64, activation='relu'))
model.add(layers.Dense(10, activation='softmax'))
model.summary()

# 编译
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
# 训练模型
history = model.fit(x_train,
                    y_train,
                    epochs=5,
                    batch_size=64,
                    validation_data=(x_test, y_test))
np.save("cnn.npy", history.history)
